AI Optimization Stack
同一次大模型训练或推理,可能同时涉及模型结构、并行策略、请求调度、计算图、Kernel、通信库和加速器。它们最终共同决定成本与性能,因此都能被称为“AI 模型优化”;但如果不先分清层次,很容易得到几个看似合理、实际上不可验证的结论:
- Kernel 快了,模型就一定同比例变快。
- 用了 GPU、NPU 或 RDMA,就是软硬协同。
- 训练优化只看算力,推理优化只看带宽。
- 框架机制可以不考虑具体硬件。
更可靠的提问方式是:
- 对象是什么? 是模型语义、参数与梯度、请求与 KV Cache、计算图、张量 Tile、通信消息,还是硬件资源?
- 物理量变了什么? 是 FLOPs、HBM 读写、峰值显存、网络字节、Kernel 启动、排队时间,还是设备能力?
- 用什么验收? 是精度、收敛、tokens/s、TTFT、TPOT、P99、单算子时延,还是能效?
七个优化层次
这里将广义的 AI 模型优化拆成七层。层与层之间不是互斥的组织边界,而是定位问题时的主要责任边界:一个改动可以跨层实现,但必须能说明每一层分别改了什么。
| 层次 | 主要控制对象 | 常见物理瓶颈 | 主要验收指标 |
|---|---|---|---|
| 模型与算法 | 目标函数、网络结构、数据与数值语义 | 有效计算量、样本效率、表达能力 | 质量、收敛速度、鲁棒性 |
| 训练系统 | 参数、梯度、优化器状态、激活、并行 Rank | 设备利用率、显存峰值、同步等待 | tokens/s、step time、可训练规模 |
| 推理服务 | 请求、Batch、KV Cache、Prefill 与 Decode | 排队、KV 容量、尾时延、空泡 | TTFT、TPOT/ITL、吞吐、P99、成本 |
| 编译与运行时 | 计算图、形状特化、内存计划、执行流 | Graph Break、重编译、调度与启动开销 | 编译时间、命中率、端到端时延 |
| 算子与 Kernel | 单个算子、Tile、Layout、融合边界 | HBM 流量、算术强度、并行度 | 单算子时延、带宽或算力利用率 |
| 通信与系统软件 | Collective、Buffer、Transport、Topology | 链路带宽、消息时延、同步与拥塞 | 有效带宽、Collective 时延、Overlap |
| 硬件 | 计算单元、存储层级、互连与指令能力 | 峰值能力、容量、功耗、拓扑约束 | 峰值规格、能效、可实现上限 |
模型与算法
模型与算法层决定“算什么”。它可以改变模型结构、训练目标、数据使用方式和数值表示,也可能改变最终可达到的质量。例如,用稀疏结构降低每个 Token 激活的参数量,或用蒸馏将教师模型的能力迁移到更小模型,首先都是算法层问题。
这一层最容易出现不公平比较。若新方法同时用了更多数据、更长训练时间或不同质量阈值,就不能只比较最终吞吐或参数量。至少要固定任务、数据预算、计算预算与质量门槛中的关键条件。
训练系统
训练系统层决定“怎么把前向、反向和参数更新组织起来”。它管理参数、梯度、优化器状态和激活的切分与生命周期,并安排数据并行、张量并行、流水线并行、上下文并行或专家并行之间的计算与通信。
以 PyTorch FSDP 的 FULL_SHARD 为例,参数在计算前按需 AllGather,梯度在反向后 ReduceScatter,并在空闲阶段重新分片;优化对象不是注意力公式,而是分布式状态的驻留方式和通信时机。[^fsdp]
训练系统优化通常关心:
- 可训练规模: 模型与状态是否能放进给定设备;
- 稳态吞吐: tokens/s、samples/s 或 step time;
- 扩展效率: 增加设备后,有效吞吐是否接近理想比例;
- 语义等价: 换并行策略、重计算或混合精度后,损失曲线和最终质量是否仍在可接受范围内。
训练并非天然“计算受限”。小 Batch、细粒度流水线、稀疏专家路由或跨节点并行都可能让通信、内存容量和调度空泡成为主导瓶颈。
推理服务
推理服务层决定“请求如何进入设备并持续生成 Token”。它没有训练中的反向和优化器状态,却多了请求到达、排队、Prefill、Decode、抢占和 KV Cache 生命周期。
其核心不是固定 Batch 的一次前向,而是在不断变化的请求集合中平衡:
- 首 Token 时延 TTFT: 用户多久看到第一次响应;
- Token 间时延 TPOT/ITL: 连续生成是否流畅;
- 吞吐: 单位时间完成多少请求或生成多少 Token;
- 尾时延与成本: P99 是否满足 SLO,每百万 Token 需要多少资源。
vLLM 的调度配置就同时涉及 Token Budget、KV Cache 容量、Prefill/Decode 策略和异步调度;官方指标也把 TTFT、Inter-Token Latency、队列和吞吐分开暴露。[^vllm-scheduler] [^vllm-metrics]
| 维度 | 训练系统 | 推理服务 |
|---|---|---|
| 主要阶段 | Forward、Backward、Optimizer Step | Prefill、Decode |
| 长驻状态 | 参数、梯度、优化器、激活 | 参数、KV Cache、请求状态 |
| 调度核心 | 并行 Rank、Micro-batch、通信依赖 | 请求、Token Budget、抢占与 KV 配额 |
| 主要指标 | step time、tokens/s、扩展效率、峰值显存 | TTFT、TPOT/ITL、吞吐、P99、成本 |
| 常见风险 | 收敛漂移、同步等待、流水线空泡 | 头阻塞、尾时延、KV 碎片或容量耗尽 |
因此,训练框架与推理框架都在提高设备利用率,但它们控制的状态机和目标函数并不相同。
编译与运行时
编译与运行时层决定“高层程序如何变成可执行图和执行序列”。典型工作包括图捕获、算子分解、形状特化、融合、内存规划、代码生成和执行流调度。
例如 torch.compile 会捕获可编译区域,交给后端生成并缓存执行产物;Guard 失败、Graph Break 或动态形状变化可能触发回退或重新编译。[^torch-compile] 这类问题不能简单归为“某个 Kernel 慢”,因为瓶颈可能发生在图边界、编译缓存或运行时调度上。
验收时既要看稳态执行,也要看:
- 首次编译成本与冷启动;
- 不同 Shape 下的缓存命中和重编译次数;
- 图覆盖率、Graph Break 与回退路径;
- 编译后峰值内存和端到端时延。
算子与 Kernel
算子与 Kernel 层决定“一个局部计算如何映射到设备”。这里会调整 Tile 大小、线程或任务分工、内存布局、流水化、向量化与融合边界。
同一个矩阵乘可以是计算受限,也可以是内存受限。判断取决于算术强度、矩阵形状和具体硬件的 Roofline,而不是“矩阵乘一定吃算力”这样的标签。NVIDIA 的矩阵乘性能指南也把形状、Tile 和算术强度作为判断性能上限的关键因素。[^matmul-guide]
算子优化必须先通过精度和边界条件测试,再比较单算子性能;最后还要回到端到端关键路径。若该算子只占总时间的 5%,即使局部快 2 倍,理想端到端收益也不到 3%:
$$
S_{\text{end-to-end}}
\frac{1}
{(1-f)+\frac{f}{S_{\text{kernel}}}}
$$
其中 $f$ 是原始运行时间中该算子的占比,$S_{\text{kernel}}$ 是单算子加速比。
通信与系统软件
通信与系统软件层决定“数据如何跨设备和节点移动”。它控制 Collective 算法、Buffer、Transport、拓扑路径、完成通知和通信流。
NCCL 提供 AllReduce、AllGather、ReduceScatter 等拓扑感知 Collective,并面向多 GPU、PCIe、NVLink 和网络互连优化通信。[^nccl] 但通信库不是完整训练框架:它不知道损失函数、优化器状态或请求 SLO,框架需要决定何时、对谁、以多大粒度调用它。
这一层的验收不能只看链路标称带宽,还应观察有效带宽、小消息时延、拓扑差异、拥塞、同步等待和计算通信 Overlap。通信变快但关键路径不变,端到端也可能没有收益。
硬件
硬件层决定“软件可用的物理能力与约束”,包括矩阵或向量计算单元、HBM、片上 SRAM、互连、DMA Engine、指令集和功耗上限。
硬件规格给出的是上限,不是应用性能。软件只有选择合适的数据类型、布局、并行度和数据搬运方式,才能逼近峰值。CUDA 编程指南对线程层次、存储层次和异步执行的描述,本质上就是软件映射硬件资源时需要遵守的执行模型。[^cuda-guide]
软硬协同
软硬协同不是与七层并列的第八层,而是一种跨层优化范式。 它要求软件设计显式利用硬件特性,或让后续硬件能力根据工作负载共同演进。
可以把它看成一条连续谱:
- 硬件感知的软件优化: 硬件不变,软件根据片上存储、Tensor Core/Cube、DMA Engine 或网络拓扑重新选择 Tile、Layout 和流水线。
- 接口协同: 编译器、运行时与通信库围绕硬件能力建立稳定抽象,使上层框架能可靠使用异步拷贝、低精度计算或设备直连。
- 共同设计: 软件工作负载反过来影响指令、缓存、互连或加速单元设计,软件也同步改变数据流和算法表达。
因此,FlashAttention 可以称为典型的硬件感知算法与 Kernel 设计:它没有修改注意力的数学结果,也没有要求先制造一种新 GPU,而是根据 GPU 的 HBM 与片上 SRAM 层级重新安排计算。只有当硬件本身也因这种数据流而发生设计变化时,才是更严格意义上的共同设计。
现实中的“软硬协同优化岗位”往往跨越训练或推理框架、编译器、Kernel、通信库和硬件性能分析。岗位名称描述的是协作半径,不是单一的软件层。
FlashAttention 机制卡
FlashAttention 适合用来观察跨层优化:数学上仍是精确注意力;算法上重新排列计算;Kernel 上利用片上存储分块;框架或编译器再负责选择和调用具体后端。下面以 FlashAttention-2 论文中的前向 Algorithm 1 与附录反向推导为依据,建立一张教学机制卡。[^flashattention2]
起因与直觉
标准缩放点积注意力为:
$$
S = \frac{QK^\mathsf{T}}{\sqrt d} + M,\qquad
P = \operatorname{softmax}(S),\qquad
O = PV
$$
其中:
- $Q,K,V\in\mathbb{R}^{B\times H\times N\times d}$;
- $B$ 是 Batch,$H$ 是注意力头数,$N$ 是序列长度,$d$ 是每个头的维度;
- $M$ 是广播到 $[B,H,N,N]$ 的 Mask;
- 语义上的 $S,P\in\mathbb{R}^{B\times H\times N\times N}$,输出 $O\in\mathbb{R}^{B\times H\times N\times d}$。
朴素实现若把完整 $S$ 和 $P$ 写回 HBM,会产生随 $N^2$ 增长的中间读写。FlashAttention 的核心直觉不是“少算一个近似注意力”,而是:
把 $Q,K,V$ 分块搬到片上快存储,在块内完成得分、Softmax 和输出累加,只把在线归一化状态与最终输出写回 HBM。
Softmax 不能简单按块独立计算,因为后来的块可能出现更大值。对第 $i$ 个 Query 行块,算法维护:
- $m_i$:目前见过的行最大值;
- $\ell_i$:在统一最大值基准下的指数和;
- $O_i$:已经归一化的部分输出。
当新块的局部最大值 $\widetilde m_i$ 到来时,旧结果通过
$\exp(m_i-m_i^{\text{new}})$ 重缩放,新块通过
$\exp(\widetilde m_i-m_i^{\text{new}})$ 对齐,从而得到与一次性 Softmax 相同的结果。
对象与生命周期
| 对象 | 逻辑形状 | 生命周期 | 是否完整写入 HBM |
|---|---|---|---|
| $Q,K,V$ | $[B,H,N,d]$ | 整个算子 | 是,作为输入 |
| $Q_i$ | $[B,H,B_r,d]$ | 一个 Query Tile | 否,片上暂存 |
| $K_j,V_j$ | $[B,H,B_c,d]$ | 一个 Key/Value Tile | 否,片上暂存 |
| $S_{ij}$ | $[B,H,B_r,B_c]$ | 一次 Tile 计算 | 否,用后释放 |
| $\widetilde P_{ij}$ | $[B,H,B_r,B_c]$ | 一次 Tile 计算 | 否,用后释放 |
| $m_i,\ell_i$ | $[B,H,B_r]$ | 跨 Key/Value Tile | 是,小状态 |
| $O_i$ | $[B,H,B_r,d]$ | 跨 Key/Value Tile | 是,滚动更新 |
| $D_i=\operatorname{rowsum}(dO_i\odot O_i)$ | $[B,H,B_r]$ | 反向的一行统计 | 可预计算保存 |
| $dQ,dK,dV$ | 同对应输入 | 整个反向 | 是,最终梯度 |
这里最重要的区分是:完整 $S$ 和 $P$ 在数学语义中存在,但在物理执行中不必被完整物化。 这也是分析编译器与 Kernel 时常用的“语义对象—物理对象”分离。
完整伪代码
下面的伪代码刻意写出 Mask、在线归一化、反向重计算与临时对象释放,便于检查输入、输出和生命周期。真实实现还会根据设备加入线程映射、流水线、向量化和融合。
1 | def flash_attention_forward(Q, K, V, mask, Br, Bc): |
反向不保存完整 $P$,而是用 $Q_i,K_j,m_i,\ell_i$ 重建当前 Tile 的概率,再累计梯度:
1 | def flash_attention_backward(Q, K, V, O, dO, mask, m, l, Br, Bc): |
这段伪代码给出了机制闭环,但不是性能实现:真实 Kernel 会重新选择循环顺序、让多个 Query Tile 共享已加载的 Key/Value Tile,并处理因果 Mask、Dropout、GQA/MQA、变长序列和数值精度等边界。
五视图
五张视图回答的是不同问题:
- 物理结构图: 哪些对象在 HBM,哪些 Tile 只在片上短暂存在;
- 因果逻辑图: 为什么减少 HBM 往返可能比单纯减少 FLOPs 更有效;
- 过程图: 前向如何在线归一化,反向如何重计算概率;
- 时序图: 临时对象何时创建、使用和释放;
- 数据流图: 输入、滚动状态、输出与梯度之间如何传递。
把它们混成一张普通“流程图”,通常会隐藏最关键的存储位置和生命周期。
论文证据与边界
论文的概览图把存储层次与端到端训练结果放在一起,说明机制收益可以穿过 Kernel 边界进入模型训练;但右侧柱状图仍然只是特定 GPT-2 配置、软件版本与 A100 环境下的一次结果,不能直接外推到 NPU、不同 Mask、超长序列或推理服务。
这个对比特别重要,因为 FlashAttention 在示例中 FLOPs 反而更多,但 HBM 读写更少、执行时间更短。它反驳了“优化就是减少计算量”的单一视角,也给出完整的因果链:
$$
\text{分块与在线 Softmax}
\rightarrow
\text{不物化完整 }S,P
\rightarrow
\text{HBM 读写下降}
\rightarrow
\text{受 IO 限制的 Shape 加速}
$$
这条链的最后一步是条件结论。若 Shape 很小、Kernel 并行度不足、实现存在额外同步,或者运行已由其他算子主导,收益会缩小。
显存影响契约
分析“省显存”时,不能只写一句“复杂度从 $O(N^2)$ 降为 $O(N)$”。端到端峰值更接近以下账本:
$$
\begin{aligned}
M_{\text{peak}}={}&
M_{\text{model state}}
+M_{\text{saved activations}}
+M_{\text{operator workspace}}\
&+M_{\text{communication buffers}}
+M_{\text{live inputs/outputs}}\
&+M_{\text{allocator reserved unused}}
+M_{\text{runtime margin}}
\end{aligned}
$$
FlashAttention 直接改变的主要是:
- 算子工作区: 不再长期物化完整 $S$ 与 $P$;
- 保存激活: 可以保存较小的归一化统计,在反向重计算局部概率;
- 生命周期: $S_{ij}$、$P_{ij}$、$dS_{ij}$ 只在当前 Tile 内短暂存在。
它不会直接减少模型参数、优化器状态、KV Cache、通信 Buffer 或内存分配器碎片。因此,算子级 $O(N^2)$ 中间量消失,不等于整个模型峰值严格按相同比例下降。
一个可执行的效果契约应写成:
- 目标项: $M_{\text{operator workspace}}$ 与相关 $M_{\text{saved activations}}$,以及 HBM 读写量。
- 动作: 通过分块、在线归一化和反向重计算,减少完整中间张量的物化并缩短生命周期。
- 代价: 额外 FLOPs、Tile 工作区、Kernel Shape 约束,以及 Mask、Dropout、精度和反向实现复杂度。
- 同场测量: 固定模型、Batch、序列长度、头维、精度、Mask/Dropout、硬件、软件版本和分配器配置,再比较单算子与端到端结果。
如何定位优化层
面对一个新的“性能优化”需求,可以按下面的顺序工作。
固定语义与指标
先写下不可变化的条件。训练任务可能要求损失曲线与最终质量等价;在线推理可能要求输出一致、P99 小于阈值;算子替换则需要明确误差容限。
然后只选一个主目标,例如:
- 在质量不下降的条件下提高训练 tokens/s;
- 在 P99 不超过 2 s 的条件下降低单 Token 成本;
- 在固定 Shape 与精度下减少注意力 HBM 流量。
没有语义约束和主指标,后面的“加速”很容易只是换了任务。
建立关键路径与内存账本
使用 Profiler 把时间拆成计算、通信、排队、编译、Kernel Launch 和空闲;把峰值内存拆成模型状态、激活、工作区、通信 Buffer、活跃输入输出、分配器保留与运行时余量。
选择最窄责任层
找到占主导的对象和生命周期后,选择能够直接控制它的最窄层次:
- 质量或有效计算量不合适,回到模型与算法层;
- 参数、梯度或激活驻留不合适,优先看训练系统;
- 排队、KV 配额或 Prefill/Decode 竞争,优先看推理服务;
- Graph Break、重编译或内存计划,优先看编译与运行时;
- 局部 HBM、Tile 或融合问题,优先看 Kernel;
- Collective、链路或拓扑等待,优先看通信;
- 现有能力无法表达目标数据流,再考虑硬件变化。
“最窄”不意味着只改一层,而是先建立清晰的主因,再检查上下游接口。
做两级验证
局部优化至少需要两级实验:
- Microbenchmark: 验证被控制的局部物理量确实改变,例如 Kernel 时延、HBM 字节或 Collective 带宽。
- End-to-end: 验证改动进入了真实关键路径,没有被调度、编译、通信或其他算子抵消。
可以把流程写成下面的通用伪代码:
1 | baseline = measure(workload, fixed_semantics, fixed_environment) |
最后报告条件、代价和未覆盖边界,而不是只报告最快的一次数字。
常见误区
- 框架优化就是算子优化: 框架可能调用更快的 Kernel,但它还负责状态、调度和跨算子依赖。反过来,单个 Kernel 可以在完全不理解训练并行或请求 SLO 的情况下被独立优化。二者通过接口连接,不应在归因时合并。
- Kernel 快等于模型快: 端到端收益受该 Kernel 在关键路径中的占比限制,还可能被 CPU 调度、编译、通信或同步抵消。必须同时报告 Microbenchmark 和 End-to-end。
- 训练吃算力,推理吃带宽: 这只能作为部分场景的经验起点。训练也可能受通信、显存或流水线空泡限制;长 Prompt 的 Prefill 可能具有较高算术强度,Decode 才更常受权重与 KV 读写影响。应针对具体阶段与 Shape 判断。
- 使用硬件 API 就是软硬协同: 调用 CUDA、CANN、NCCL 或 HCCL 只说明使用了某个接口。只有软件设计显式围绕硬件约束重构数据流,或软硬件共同改变接口与能力,才形成有意义的协同证据。
- 换更强硬件就不需要软件优化: 新硬件会改变算力、带宽、容量和拓扑的比例,也会移动瓶颈。旧软件的数据布局、并行策略或调度可能无法利用新能力,甚至让更高峰值长期闲置。
总结
AI 模型优化不是一条从算法到芯片的单向流水线,而是七个相互反馈的控制层:
- 模型与算法决定算什么;
- 训练系统组织状态、前后向和分布式并行;
- 推理服务组织请求、Token 与 KV Cache;
- 编译与运行时把程序变成稳定的执行图;
- 算子与 Kernel把局部计算映射到设备;
- 通信与系统软件搬运跨设备数据;
- 硬件提供物理能力与约束。
软硬协同贯穿这些层。判断一个优化属于哪里,不看它使用了多少流行名词,而看控制对象、物理变化、验收指标和证据边界是否闭环。FlashAttention 的价值也正在于此:它把数学语义、IO 复杂度、片上存储、Kernel 数据流和端到端训练收益连成了一条可验证的因果链。
可复用提示词
当需要分析新的 AI 优化方案时,可以复用下面的提示词:
1 | 请不要先按项目名称归类,而是分析这个优化的: |
参考资料
[^fsdp]: PyTorch, Fully Sharded Data Parallel.
[^vllm-scheduler]: vLLM, Scheduler Configuration.
[^vllm-metrics]: vLLM, Metrics.
[^torch-compile]: PyTorch, torch.compile.
[^matmul-guide]: NVIDIA, Matrix Multiplication Background User’s Guide.
[^nccl]: NVIDIA, NCCL Documentation.
[^cuda-guide]: NVIDIA, CUDA Programming Guide.
[^flashattention2]: Tri Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, arXiv:2205.14135v2, 2023.